# coding=utf-8
import os
import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from math import log
import os.path
from chains.C15_bubble_title import Bubble
from encode.text_encoder import TextToVec
import numpy as np
import aiomysql
import faiss
import time
import ast
from config.config import config

from utils.db_pool import LawDB
from utils.logger import logger

print('开始')
time0 = time.time()
text2vec = TextToVec(os.path.join(config.MODEL_DIR, 'bge-large-zh-v1_5_'))

time1 = time.time()
print('模型加载完毕', time1 - time0)

index = faiss.read_index(os.path.join(config.VECTOR_DIR, 'IPO_map2_id_content_dict_bge_large_zh_v1_5.faiss'))
time2 = time.time()
print("faiss_model加载时间：", time2 - time1)


# 根据 板块、行业、地域 召回申报材料
async def recall_material_list(content_text: str, k: int, sector: str, industry: str, region: str, problem_type_list: list):
    # 设置 板块、行业 报错， 都能不输入

    '''
    :param content_text: str 招股书切片
    :param k: top k
    :param sector: 板块 （主板, 科创板）
    :param industry: 行业
    :param region: 地域
    :param problem_type_list: 问题分类标签 list[str]
    :return: material_id_list 召回的 材料ID 的二维列表
    '''

    s_time = time.time()
    # sector_list = ['主板', '科创板', '创业板']
    sector_list = ['沪市主板', '深市主板', '科创板', '创业板']
    industry_list = ['农业', '林业', '畜牧业', '渔业', '农、林、牧、渔服务业', '煤炭开采和洗选业', '石油和天然气开采业',
                 '黑色金属矿采选业', '有色金属矿采选业', '非金属矿采选业',
                 '开采辅助活动', '其他采矿业', '农副食品加工业', '食品制造业', '酒、饮料和精制茶制造业', '烟草制品业',
                 '纺织业', '纺织服装、服饰业',
                 '皮革、毛皮、羽毛及其制品和制鞋业',
                 '木材加工和木、竹、藤、棕、草制品业', '家具制造业', '造纸和纸制品业', '印刷和记录媒介复制业',
                 '文教、工美、体育和娱乐用品制造业', '石油加工、炼焦和核燃料加工业',
                 '化学原料和化学制品制造业', '医药制造业', '化学纤维制造业', '橡胶和塑料制品业', '非金属矿物制品业',
                 '黑色金属冶炼和压延加工业', '有色金属冶炼和压延加工业',
                 '金属制品业',
                 '通用设备制造业', '专用设备制造业', '汽车制造业', '铁路、船舶、航空航天和其他运输设备制造业',
                 '电气机械和器材制造业', '计算机、通信和其他电子设备制造业', '仪器仪表制造业',
                 '其他制造业', '废弃资源综合利用业', '金属制品、机械和设备修理业', '电力、热力生产和供应业',
                 '燃气生产和供应业', '水的生产和供应业', '房屋建筑业', '土木工程建筑业',
                 '建筑安装业', '建筑装饰和其他建筑业', '批发业', '零售业', '铁路运输业', '道路运输业', '水上运输业',
                 '航空运输业', '管道运输业', '装卸搬运和其他运输代理业',
                 '装卸搬运和运输代理业', '仓储业', '邮政业', '住宿业', '餐饮业', '电信、广播电视和卫星传输服务',
                 '互联网和相关服务', '软件和信息技术服务业', '货币金融服务',
                 '资本市场服务', '保险业', '其他金融业', '房地产业', '租赁业', '商务服务业', '研究和试验发展',
                 '专业技术服务业', '科技推广和应用服务业', '水利管理业',
                 '生态保护和环境治理业', '公共设施管理业', '居民服务业', '机动车、电子产品和日用产品修理业',
                 '其他服务业', '教育', '卫生', '社会工作', '新闻和出版业',
                 '广播、电视、电影和影视录音制作业', '文化艺术业', '体育', '娱乐业', '综合']


    if sector not in sector_list:
        raise ValueError('请输入正确的 板块')
    elif industry not in industry_list:
        raise ValueError('请输入正确的 行业')

    conn = None
    cursor = None
    try:
        # conn = await aiomysql.connect(
        #     host='nj-cdb-1mwvwkzt.sql.tencentcdb.com',
        #     port=63972,
        #     user='root',
        #     password='grow1234',
        #     charset='utf8mb4',
        #     db='stock_exchange'
        # )

        # conn = await aiomysql.connect(
        #     host='localhost',
        #     port=3306,
        #     user='root',
        #     password='123456',
        #     charset='utf8mb4',
        #     db='local_stock_exchange'
        # )

        conn = await aiomysql.connect(
            host='192.168.2.244',
            port=3306,
            user='root',
            password='linkAge@12345',
            charset='utf8mb4',
            db='stock_exchange'
        )

        cursor = await conn.cursor()
        conditions = []
        values = []
        if sector and sector != "" and sector in sector_list:
            # 找到相同板块【沪/深/创】的不同行业分类
            await cursor.execute("SELECT DISTINCT industry FROM home WHERE sector = %s", (sector,))
            sector_results = await cursor.fetchall()
            my_industry_list = []
            for sector_tuple in sector_results:
                my_industry_list.append(str(sector_tuple[0]))
            print('my_industry_list', my_industry_list)
            conditions.append("sector = %s")
            values.append(sector)
        else:
            await cursor.execute("SELECT DISTINCT industry FROM home")
            sector_results = await cursor.fetchall()
            my_industry_list = []
            for sector_tuple in sector_results:
                my_industry_list.append(str(sector_tuple[0]))

            conditions.append("sector IS NOT NULL")

        if industry and industry != "" and industry in my_industry_list:
            # 找到相同板块【沪/深/创】相同行业【A/B/C】不同地区分类【比如浙江】
            await cursor.execute("SELECT DISTINCT registered_address FROM home WHERE sector = %s AND industry = %s", (sector, industry,))
            region_results = await cursor.fetchall()
            my_region_list = []
            for region_tuple in region_results:
                my_region_list.append(str(region_tuple[0]))

            conditions.append("industry = %s")
            values.append(industry)
        else:
            print('该行业不在我们数据库中，使用该板块下所有行业匹配')
            # 找到相同板块【沪/深/创】的不同地区分类【比如浙江】
            await cursor.execute("SELECT DISTINCT registered_address FROM home WHERE sector = %s", (sector,))
            region_results = await cursor.fetchall()
            my_region_list = []
            for region_tuple in region_results:
                my_region_list.append(str(region_tuple[0]))

            conditions.append("industry IS NOT NULL")

        if region and region != "" and region in my_region_list:
            conditions.append("registered_address = %s")
            values.append(region)
        else:
            conditions.append("registered_address IS NOT NULL")
        # 根据板块，行业，找到home_id
        query = "SELECT home_id FROM home WHERE " + " AND ".join(conditions)
        await cursor.execute(query, tuple(values))

        home_results = await cursor.fetchall()
        home_id_list = []
        for home_tuple in home_results:
            home_id_list.append(home_tuple[0])
        print('home_id_list长度', len(home_id_list))

        if problem_type_list == ['该招股书切片不适合被询问']:
            problem_type_list = []
        print(type(problem_type_list))

        if len(problem_type_list) == 0:
            await cursor.execute("SELECT problem_id, problem_classification_v3 FROM problem WHERE project_id IN %s", (home_id_list,))
            problem_results = await cursor.fetchall()
            print('problem_results长度', len(problem_results))
            problem_id_list = []
            for problem_tuple in problem_results:
                if len(problem_tuple[1]) > 0:  # 如果问题分类 不为空
                    problem_id_list.append(problem_tuple[0])
        else:
            # 找到和他相似行业，相似板块的问题。比如，都是沪市主板，都是造纸和印刷品页
            await cursor.execute("SELECT problem_id, problem_classification_v3 FROM problem WHERE project_id IN %s", (home_id_list,))
            problem_results = await cursor.fetchall()
            print('problem_results长度', len(problem_results))
            problem_id_list = []
            problem_id_list_all = []
            for problem_tuple in problem_results:
                if len(problem_tuple[1]) > 0:  # 如果问题分类 不为空
                    problem_id_list_all.append(problem_tuple[0])
                    # 如果当前的问题分类标签 与同板块行业下的问询问题分类标签 存在交集，将问题id加入列表
                    if set(problem_type_list).intersection(set(ast.literal_eval(problem_tuple[1]))):
                        problem_id_list.append(problem_tuple[0])
            if len(problem_id_list) == 0:
                # 切片的问题分类标签，与同板块、同行业下的问题分类标签没有交集，使用同板块行业下所有问题ID
                print('切片问题分类标签 与同板块行业下的问询问题分类标签 全部没有交集，使用同板块行业下所有问题ID')
                problem_id_list = problem_id_list_all
                if len(problem_id_list_all) == 0:
                    print('使用同板块行业下所有问题ID依然问题ID列表长度为0')
                    return []
        print('problem_id_list长度', len(set(problem_id_list)))
        # await cursor.execute("SELECT 问题ID FROM problem WHERE 项目ID in %s AND 问题题目 LIKE %s", (home_id_list, f"%{problem_type}%"))
        # await cursor.execute("SELECT 问题ID FROM problem WHERE 项目ID IN %s AND 问题题目 LIKE %s", (sector, problem_type, ))

        # 找到对应问题的材料id，材料顺序
        await cursor.execute("SELECT material_id, material_order FROM material WHERE problem_id IN %s", (problem_id_list,))
        Material_results = await cursor.fetchall()
        Material_id_list = []
        for Material_tuple in Material_results:
            if Material_tuple[1] == 0:
                continue
            Material_id_list.append(Material_tuple[0])
        print('len(Material_id_list)', len(Material_id_list))
        if len(Material_id_list) == 0:
            return []

        time1 = time.time()
        query_embedding = text2vec.text2vec(content_text)
        query_embedding = np.array([query_embedding])
        query_embedding = query_embedding.astype(np.float32)
        time2 = time.time()
        print("文本编码时间：", time2 - time1)

        start_time = time.time()
        # 根据材料id，找到对应的向量
        subset_vectors = index.reconstruct_batch(Material_id_list)  # 重组
        # print(subset_vectors.shape)
        faiss_index = faiss.IndexFlatL2(1024)
        subset_index = faiss.IndexIDMap(faiss_index)
        # 添加向量，和材料id
        subset_index.add_with_ids(subset_vectors, Material_id_list)
        print('新faiss长度', subset_index.ntotal, len(Material_id_list))
        # 对query 向量进行搜索
        distances, laws_index = subset_index.search(query_embedding, len(Material_id_list))
        ID_list = laws_index[0]
        end_time = time.time()
        print("单个问题查询时间:", end_time - start_time)
        # material_id_list = []
        # for ID_list in id_list:
        material_list = []
        for ID in ID_list[:k]:
            await cursor.execute("SELECT material_id, material_content FROM material WHERE material_id = %s", (ID,))
            material_results = await cursor.fetchone()
            material_list.append({int(ID): material_results[1]})
            # print(type(list(ID_list)))
            # print('ID_list', ID_list)
            # material_id_list.append(material_list)
    except Exception as e:
        logger.error(f"C2 error: {e}", exc_info=True)
        # print(f"C2 error: {e}")
        return f"C2 error : {e}"
    finally:
        if cursor:
            await cursor.close()  # 使用 await 关键字
        if conn:
            conn.close()
    e_time = time.time()
    print('总耗时：', e_time - s_time)
    return material_list

# 根据 板块、行业、地域 召回 材料距离
async def recall_material_distances_db(pool: aiomysql.Pool, content_text: str, k: int, sector: str, industry: str, region: str, problem_type_list: list):
    # 设置 板块、行业 报错， 都能不输入

    '''
    :param content_text: str 招股书切片
    :param k: top k
    :param sector: 板块 （主板, 科创板）
    :param industry: 行业
    :param region: 地域
    :param problem_type_list: 问题分类标签 list[str]
    :return: material_id_list 召回的 材料ID 的二维列表
    '''

    s_time = time.time()
    sector_list = ['沪市主板', '深市主板', '科创板', '创业板']
    industry_list = ['农业', '林业', '畜牧业', '渔业', '农、林、牧、渔服务业', '煤炭开采和洗选业', '石油和天然气开采业',
                 '黑色金属矿采选业', '有色金属矿采选业', '非金属矿采选业',
                 '开采辅助活动', '其他采矿业', '农副食品加工业', '食品制造业', '酒、饮料和精制茶制造业', '烟草制品业',
                 '纺织业', '纺织服装、服饰业',
                 '皮革、毛皮、羽毛及其制品和制鞋业',
                 '木材加工和木、竹、藤、棕、草制品业', '家具制造业', '造纸和纸制品业', '印刷和记录媒介复制业',
                 '文教、工美、体育和娱乐用品制造业', '石油加工、炼焦和核燃料加工业',
                 '化学原料和化学制品制造业', '医药制造业', '化学纤维制造业', '橡胶和塑料制品业', '非金属矿物制品业',
                 '黑色金属冶炼和压延加工业', '有色金属冶炼和压延加工业',
                 '金属制品业',
                 '通用设备制造业', '专用设备制造业', '汽车制造业', '铁路、船舶、航空航天和其他运输设备制造业',
                 '电气机械和器材制造业', '计算机、通信和其他电子设备制造业', '仪器仪表制造业',
                 '其他制造业', '废弃资源综合利用业', '金属制品、机械和设备修理业', '电力、热力生产和供应业',
                 '燃气生产和供应业', '水的生产和供应业', '房屋建筑业', '土木工程建筑业',
                 '建筑安装业', '建筑装饰和其他建筑业', '批发业', '零售业', '铁路运输业', '道路运输业', '水上运输业',
                 '航空运输业', '管道运输业', '装卸搬运和其他运输代理业',
                 '装卸搬运和运输代理业', '仓储业', '邮政业', '住宿业', '餐饮业', '电信、广播电视和卫星传输服务',
                 '互联网和相关服务', '软件和信息技术服务业', '货币金融服务',
                 '资本市场服务', '保险业', '其他金融业', '房地产业', '租赁业', '商务服务业', '研究和试验发展',
                 '专业技术服务业', '科技推广和应用服务业', '水利管理业',
                 '生态保护和环境治理业', '公共设施管理业', '居民服务业', '机动车、电子产品和日用产品修理业',
                 '其他服务业', '教育', '卫生', '社会工作', '新闻和出版业',
                 '广播、电视、电影和影视录音制作业', '文化艺术业', '体育', '娱乐业', '综合']


    if sector not in sector_list:
        raise ValueError('请输入正确的 板块')
    elif industry not in industry_list:
        raise ValueError('请输入正确的 行业')


    try:
        if not pool:
            raise RuntimeError("数据库连接池未初始化")

        async with pool.acquire() as conn:
            async with conn.cursor() as cursor:
                conditions = []
                values = []
                if sector and sector != "" and sector in sector_list:
                    await cursor.execute("SELECT DISTINCT industry FROM home WHERE sector = %s", (sector,))
                    sector_results = await cursor.fetchall()
                    my_industry_list = []
                    for sector_tuple in sector_results:
                        my_industry_list.append(str(sector_tuple[0]))

                    conditions.append("sector = %s")
                    values.append(sector)
                else:
                    await cursor.execute("SELECT DISTINCT industry FROM home")
                    sector_results = await cursor.fetchall()
                    my_industry_list = []
                    for sector_tuple in sector_results:
                        my_industry_list.append(str(sector_tuple[0]))

                    conditions.append("sector IS NOT NULL")

                if industry and industry != "" and industry in my_industry_list:
                    await cursor.execute("SELECT DISTINCT registered_address FROM home WHERE sector = %s AND industry = %s", (sector, industry,))
                    region_results = await cursor.fetchall()
                    my_region_list = []
                    for region_tuple in region_results:
                        my_region_list.append(str(region_tuple[0]))

                    conditions.append("industry = %s")
                    values.append(industry)
                else:
                    print('该行业不在我们数据库中，使用该板块下所有行业匹配')

                    await cursor.execute("SELECT DISTINCT registered_address FROM home WHERE sector = %s", (sector,))
                    region_results = await cursor.fetchall()
                    my_region_list = []
                    for region_tuple in region_results:
                        my_region_list.append(str(region_tuple[0]))

                    conditions.append("industry IS NOT NULL")

                if region and region != "" and region in my_region_list:
                    conditions.append("registered_address = %s")
                    values.append(region)
                else:
                    conditions.append("registered_address IS NOT NULL")

                query = "SELECT home_id FROM home WHERE " + " AND ".join(conditions)
                await cursor.execute(query, tuple(values))

                home_results = await cursor.fetchall()
                home_id_list = []
                for home_tuple in home_results:
                    home_id_list.append(home_tuple[0])
                print('home_id_list长度', len(home_id_list))

                if problem_type_list == ['该招股书切片不适合被询问']:
                    problem_type_list = []
                print(type(problem_type_list))

                if len(problem_type_list) == 0:
                    await cursor.execute("SELECT problem_id, problem_classification_v3 FROM problem WHERE project_id IN %s", (home_id_list,))
                    problem_results = await cursor.fetchall()
                    print('problem_results长度', len(problem_results))
                    problem_id_list = []
                    for problem_tuple in problem_results:
                        if len(problem_tuple[1]) > 0:  # 如果问题分类 不为空
                            problem_id_list.append(problem_tuple[0])
                else:
                    await cursor.execute("SELECT problem_id, problem_classification_v3 FROM problem WHERE project_id IN %s", (home_id_list,))
                    problem_results = await cursor.fetchall()
                    print('problem_results长度', len(problem_results))
                    problem_id_list = []
                    problem_id_list_all = []
                    for problem_tuple in problem_results:
                        if len(problem_tuple[1]) > 0:  # 如果问题分类 不为空
                            problem_id_list_all.append(problem_tuple[1])
                            if set(problem_type_list).intersection(set(ast.literal_eval(problem_tuple[1]))):
                                problem_id_list.append(problem_tuple[0])
                if len(problem_id_list) == 0:
                    print('切片问题分类标签 与同板块行业下的问询问题分类标签 全部没有交集，使用同板块行业下所有问题ID')
                    problem_id_list = problem_id_list_all
                    if len(problem_id_list_all) == 0:
                        print('使用同板块行业下所有问题ID依然问题ID列表长度为0')
                        return []
                print('problem_id_list长度', len(set(problem_id_list)))
                # await cursor.execute("SELECT 问题ID FROM problem WHERE 项目ID in %s AND 问题题目 LIKE %s", (home_id_list, f"%{problem_type}%"))
                # await cursor.execute("SELECT 问题ID FROM problem WHERE 项目ID IN %s AND 问题题目 LIKE %s", (sector, problem_type, ))

                await cursor.execute("SELECT material_id, material_order FROM material WHERE problem_id IN %s", (problem_id_list,))
                Material_results = await cursor.fetchall()
                Material_id_list = []
                for Material_tuple in Material_results:
                    if Material_tuple[1] == 0:
                        continue
                    Material_id_list.append(Material_tuple[0])
                print('len(Material_id_list)', len(Material_id_list))
                if len(Material_id_list) == 0:
                    return []

                time1 = time.time()
                query_embedding = text2vec.text2vec(content_text)
                query_embedding = np.array([query_embedding])
                query_embedding = query_embedding.astype(np.float32)
                time2 = time.time()
                print("文本编码时间：", time2 - time1)

                start_time = time.time()
                subset_vectors = index.reconstruct_batch(Material_id_list)  # 重组
                # print(subset_vectors.shape)
                faiss_index = faiss.IndexFlatL2(1024)
                subset_index = faiss.IndexIDMap(faiss_index)
                subset_index.add_with_ids(subset_vectors, Material_id_list)
                print('新faiss长度', subset_index.ntotal, len(Material_id_list))

                distances, laws_index = subset_index.search(query_embedding, len(Material_id_list))
                ID_list = laws_index[0]
                Distances_list = distances[0]
                end_time = time.time()
                print("单个问题查询时间:", end_time - start_time)
                # material_id_list = []
                # for ID_list in id_list:
                material_list = []
                for num, ID in enumerate(ID_list[:k]):
                    await cursor.execute("SELECT material_id, material_content FROM material WHERE material_id = %s", (ID,))
                    material_results = await cursor.fetchone()
                    material_list.append({int(ID): {"材料内容": material_results[1], "材料距离": Distances_list[num]}})
    except Exception as e:
        print(f"C2 error: {e}")
        # logger错误的堆栈信息
        logger.error(f"C2 error: {e}", exc_info=True)
        return f"C2 error : {e}"
    # finally:
    #     if pool:
    #         pool.close()
    #         await pool.wait_closed()
    #         print("数据库连接池已关闭")

    e_time = time.time()
    print('总耗时：', e_time - s_time)
    return material_list


if __name__ == '__main__':
    import asyncio

    # content_text_list = ["内容1", "内容2"]
    content_text_list = [
        "发行人与齐齐哈尔博实畜牧服务有限公司就支持和牛养殖事宜签署了《支持和牛养殖协议》，约定发行人为齐齐哈尔博实畜牧服务有限公司提供6,000万元额度的专项资金支持，用于和牛养殖服务；齐齐哈尔博实畜牧服务有限公司以其拥有的牛只对上述专项资金支持额度提供抵押担保，齐齐哈尔博实畜牧服务有限公司的直接股东彭辉及关联方龙江县彭辉肉牛养殖场、龙江黑牛牧业有限公司梅里斯分公司提供连带责任保证。",
        "发行人的鸡肉和猪肉等原料肉采购通常是与客户先签署采购框架合同采购的具体产品、单价和数量后续由双方签订的书面订单或即时通讯等方式确认的电子订单为准。"]
    content_text_list = '''发行人将部分非核心工序如机械加工、表面处理等环节的工作交付由委托加工商完成，满足订单需求的同时，集中优势资源于产品的核心技术环节和关键工序，提高生产效率，最大化发挥发行人方案设计、工艺解决方案的核心竞争优势。'''
    # material_id_list = asyncio.run(recall_material(content_text_list, 10))

    #######################################################################################################

    content_text = "2019 年纸浆的价格有所下滑，2020 年整体波动较大，2021 年全球原材料 市场价格普遍上涨带动纸浆价格相比 2020 年有较大涨幅。如果未来国际国内纸 浆市场价格波动频繁且幅度加大，公司无法通过扩大销售规模摊薄单位固定成 本及采取其他有效的成本控制措施，不能将材料成本上升及时转嫁给客户或者 通过提升产品附加值以提高产品价格，则会影响公司的盈利水平。"
    sector = '沪市主板'
    industry = '造纸和纸制品业'
    region = None
    problem_type_list = ["风险提示与说明"]
    material_list = asyncio.run(recall_material_list(content_text, 3, sector, industry, region, problem_type_list))
    print('material_list:')
    print(material_list)
    # print(*material_list, sep='\n')
    # 仪器仪表制造业

    # print('material_id_list:')
    # print(material_id_list)
    # print(*material_id_list[0], sep='\n\n')
    # print(type(material_id_list))  # [[1720, 1805, 297], [147, 70, 2555]]
